import os
import cv2
import numpy as np


def load_image_tile(image_dir, file_name):
    image_path = os.path.join(image_dir, file_name)
    image = cv2.imread(image_path)
    return image


def load_flow_tile(flow_dir, row, col):
    flow_path = os.path.join(flow_dir, f"tr{row}-tc{col}.npy")
    flow = np.load(flow_path)
    return flow


def extend_image_and_flow(image, flow, border_size=512):
    # 扩展图像
    extended_image = cv2.copyMakeBorder(image, border_size, border_size, border_size, border_size, cv2.BORDER_CONSTANT,
                                        value=[0, 0, 0])

    # 扩展光流
    flow_extended = cv2.copyMakeBorder(flow, border_size, border_size, border_size, border_size, cv2.BORDER_CONSTANT,
                                       value=0)

    # 填充上边界
    flow_extended[:border_size, border_size:-border_size, :] = flow[0, :, :]
    # 填充下边界
    flow_extended[-border_size:, border_size:-border_size, :] = flow[-1, :, :]
    # 填充左边界
    flow_extended[border_size:-border_size, :border_size, :] = flow[:, 0, np.newaxis]
    # 填充右边界
    flow_extended[border_size:-border_size, -border_size:, :] = flow[:, -1, np.newaxis]

    # 填充四个角
    flow_extended[:border_size, :border_size, :] = flow[0, 0, :]
    flow_extended[-border_size:, :border_size, :] = flow[-1, 0, :]
    flow_extended[:border_size, -border_size:, :] = flow[0, -1, :]
    flow_extended[-border_size:, -border_size:, :] = flow[-1, -1, :]

    return extended_image, flow_extended


def apply_flow_to_tile(image, flow):
    h, w = image.shape[:2]
    flow_map = np.zeros((h, w, 2), dtype=np.float32)
    flow_map[..., 0] = flow[..., 0] + np.arange(w)
    flow_map[..., 1] = flow[..., 1] + np.arange(h)[:, np.newaxis]

    remapped_image = cv2.remap(image, flow_map, None, cv2.INTER_LINEAR)
    return remapped_image


def save_tile(image, output_dir, file_name):
    output_path = os.path.join(output_dir, file_name)
    cv2.imwrite(output_path, image)


def process_large_image(image_dir, flow_dir, output_dir, tile_size=1024, border_size=512):
    os.makedirs(output_dir, exist_ok=True)

    # 遍历所有图像块文件
    for image_file in os.listdir(image_dir):
        if image_file.endswith('.png'):
            # 提取行列编号
            parts = image_file.split('tr')[1].split('-tc')
            row = int(parts[0])
            col = int(parts[1].split('.')[0])

            # 加载当前块及其光流场
            image = load_image_tile(image_dir, image_file)
            flow = load_flow_tile(flow_dir, row, col)

            # 扩展边界
            extended_image, extended_flow = extend_image_and_flow(image, flow, border_size)

            # 应用光流场
            transformed_tile = apply_flow_to_tile(extended_image, extended_flow)

            # 保存扩展后的图像
            save_tile(transformed_tile, output_dir, image_file)


# 示例用法
image_dir = 'E:\\wafer52\\11867_32nm_big_combine'  # 替换为实际图像块目录
flow_dir = 'E:\\wafer52\\32nm_flow_big_combine'  # 替换为实际光流场目录
output_dir = 'E:\\wafer52\\11867_32nm_big_corse_align'  # 替换为实际输出目录

process_large_image(image_dir, flow_dir, output_dir)
